package com.yj.rocketmq.listener;

import com.yj.rocketmq.RocketMqConstants;
import com.yj.rocketmq.api.entity.MessageRecord;
import com.yj.rocketmq.api.mapper.MessageRecordMapper;
import lombok.extern.slf4j.Slf4j;
import org.apache.rocketmq.client.producer.LocalTransactionState;
import org.apache.rocketmq.client.producer.TransactionListener;
import org.apache.rocketmq.common.message.Message;
import org.apache.rocketmq.common.message.MessageExt;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.util.concurrent.ConcurrentHashMap;


/**
 * @author 杨旭平
 * @date 19-6-20
 */
@Slf4j
@Component
public class LocalTransactionListener implements TransactionListener {

    private ConcurrentHashMap<String, Integer> countHashMap = new ConcurrentHashMap<>();

    private final static int MAX_COUNT = 2;

    @Autowired
    private MessageRecordMapper messageRecordMapper;

    /**
     * 执行本地事物校验
     *
     * @param message
     * @param arg
     * @return
     */
    @Override
    public LocalTransactionState executeLocalTransaction(Message message, Object arg) {
        log.info("{}topic={}，tag={}，消息体={},arg={}", RocketMqConstants.LOG_TRANSACTION_PREFIX,
                message.getTopic(), message.getTags(), new String(message.getBody()), arg);
        return LocalTransactionState.UNKNOW;
    }

    /**
     * 本地事物校验
     *
     * @param messageExt
     * @return
     */
    @Override
    public LocalTransactionState checkLocalTransaction(MessageExt messageExt) {
        //获取消息id
        String msgId = messageExt.getProperty("KEYS");
        if (this.ifExistMessage(msgId)) {
            return LocalTransactionState.COMMIT_MESSAGE;
        }
        return rollBackOrUnknown(msgId);
    }


    /**
     * 重试,决定是否回滚
     *
     * @param msgId
     * @return
     */
    private LocalTransactionState rollBackOrUnknown(String msgId) {
        Integer num = countHashMap.get(msgId);

        if (num != null && ++num > MAX_COUNT) {
            countHashMap.remove(msgId);
            return LocalTransactionState.ROLLBACK_MESSAGE;
        }

        if (num == null) {
            num = 1;
        }

        countHashMap.put(msgId, num);
        return LocalTransactionState.UNKNOW;
    }

    /**
     * 校验事物消息是否存在
     * 用来校验数据库层面是否已经操作就绪,差commit
     *
     * @param msgId
     * @return
     */
    private Boolean ifExistMessage(String msgId) {
        MessageRecord messageRecord = messageRecordMapper.selectById(msgId);
        if (null != messageRecord) {
            return true;
        } else {
            return false;
        }
    }
}
